import streamlit as st


def redbook_page():
    # import streamlit as st

    from api.redbook_api import get_redbook_content

    # st.set_page_config(layout="wide")

    st.title("📕小红书文案生成器")

    # 定义布局（侧边栏）
    with st.sidebar:
        #  定义标题
        st.divider()
        # 通义⼤模型密钥输⼊
        key = st.text_input("请输⼊通义大模型api key：", type="password")
        # key获取跳转链接
        st.page_link(
            "https://bailian.console.aliyun.com/?spm=5176.28326591.0.0.40f76ee1vUR5Of&accounttraceid=d873c9fe82ea443d9f3170f30dd5690cfnng#/api_key",
            label="⽆key？请点我获取", icon="♿")

        st.divider()

    st.header("📕小红书文案成器")
    subject = st.text_input("请输入文章主题")
    temperature = st.slider("创造性", min_value=0.0, max_value=2.0, value=1.0, step=0.1)
    flag = st.button("生成")
    if flag:
        if not key:
            st.info("请先输入通义大模型api key")
            st.stop()
        if not subject:
            st.info("请输入文章主题")
            st.stop()

        # Create a list to store the results
        results = []

        # Create a placeholder for the table
        table_placeholder = st.empty()

        # Generate 6 sets of content
        for _ in range(6):
            with st.spinner("正在生成文案，请稍等🥵🥵🥵"):
                title, response = get_redbook_content(key, subject, temperature)
                results.append({"标题": title, "正文": response})

                # Update the table
                table_placeholder.table(results)


def video_page():
    # import streamlit as st

    from video.video_demo import get_video_content

    # st.set_page_config(layout="wide")

    st.title("🎦视频文案生成器")

    # 定义布局（侧边栏）
    with st.sidebar:
        #  定义标题
        st.title("视频文案生成器")
        st.divider()
        # 通义⼤模型密钥输⼊
        key = st.text_input("请输⼊通义大模型api key：", type="password")
        # key获取跳转链接
        st.page_link(
            "https://bailian.console.aliyun.com/?spm=5176.28326591.0.0.40f76ee1vUR5Of&accounttraceid=d873c9fe82ea443d9f3170f30dd5690cfnng#/api_key",
            label="⽆key？请点我获取", icon="♿")

        st.divider()

    st.header("🎦视频文案成器")
    subject = st.text_input("请输入视频主题")
    time = st.number_input("请输入视频时长（单位：分钟）", min_value=0.0, max_value=6.0, value=0.1)
    temperature = st.slider("创造性", min_value=0.0, max_value=2.0, value=1.0, step=0.1)
    flag = st.button("生成")
    if flag:
        if not key:
            st.info("请先输入通义大模型api key")
            st.stop()
        if not subject:
            st.info("请输入视频主题")
            st.stop()

        with st.spinner("正在生成视频文案，请稍等🥵🥵🥵"):
            title, response = get_video_content(key, subject, time, temperature)
            st.divider()
            st.header("🎦标题：")
            st.title(title)
            st.divider()
            st.write(response)


def ai_page():
    import time

    # import streamlit as st

    from api.demo_api import Tongyi_llm
    from rag.file_rag import get_response_from_file

    # st.set_page_config(layout="wide")

    st.title("CLOSEAI")

    # 定义布局（侧边栏）
    with st.sidebar:
        #  定义标题
        st.title("CLOSEAI")
        st.divider()
        # 通义⼤模型密钥输⼊
        key = st.text_input("请输⼊通义大模型api key：", type="password")
        # key获取跳转链接
        st.page_link(
            "https://bailian.console.aliyun.com/?spm=5176.28326591.0.0.40f76ee1vUR5Of&accounttraceid=d873c9fe82ea443d9f3170f30dd5690cfnng#/api_key",
            label="⽆key？请点我获取", icon="♿")

        st.divider()

        # 在侧边栏中添加链的类型选择
        chain_choice = st.selectbox(
            '选择链的类型',
            ('使用新版create_retriever_chain', '使用老版ConversationalRetrievalChain')
        )

        # 当选择使用老版ConversationalRetrievalChain时，展示策略选择
        if chain_choice == '使用老版ConversationalRetrievalChain':
            strategy_choice = st.selectbox(
                '选择ConversationalRetrievalChain的策略',
                ('stuff', 'map_reduce', 'refine', 'map_rerank')
            )
        else:
            strategy_choice = None

        st.session_state["chain_choice"] = chain_choice
        st.session_state["strategy_choice"] = strategy_choice

    # 创建一个可折叠的内容区域
    expander = st.expander('选择模型')
    expander.write('只有通义千问')

    # 定义⼀个模拟流式输出的⽅法
    def stream_data(result_text):
        # 将空格替换为特殊字符
        result_text = result_text.replace(' ', '_SPACE_')
        for word in result_text.split(" "):
            # 在输出时将特殊字符替换回空格
            yield word.replace('_SPACE_', ' ')
            time.sleep(0.02)

    # 定义⼀个聊天记录到会话管理中
    if "messages" not in st.session_state:
        st.session_state["messages"] = [{"role": "ai", "content": "你好，我是逆蝶，你有什么想说的"}]
        st.session_state["tongyi"] = Tongyi_llm()
        st.session_state["tongyi"].create_llm(key)
        st.session_state["chain"] = st.session_state["tongyi"].get_chain()

    print(st.session_state["messages"])

    # 文件上传
    uploaded_files = st.file_uploader("选择一个或多个文件", accept_multiple_files=True, type=["txt", "pdf", "csv"])

    if uploaded_files:
        st.write('你选择了以下文件：')
        for uploaded_file in uploaded_files:
            st.write(uploaded_file.name)

    for message in st.session_state["messages"]:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    user_input = st.chat_input("我在看着你👁️")

    if user_input:
        st.session_state["messages"].append({"role": "human", "content": user_input})
        with st.chat_message("human"):
            st.write_stream(stream_data(user_input))

    if user_input:
        if not key:
            st.info("请输⼊你的密钥")
            st.stop()
        with st.spinner("AI正在思考🥵🥵🥵🥵🥵"):
            if uploaded_files:
                response, relvt = get_response_from_file(key, uploaded_files, user_input, chain_choice, strategy_choice)
                result = "请根据以下信息" + response + "回答" + user_input + "你的回答开头为：根据您提供的文件，[你的内容]"
                message = st.session_state["tongyi"].invoke(user_input=result)
            else:
                message = st.session_state["tongyi"].invoke(user_input=user_input)
            st.session_state["messages"].append({"role": "ai", "content": message})
            with st.chat_message("ai"):
                st.write_stream(stream_data(message))
            if uploaded_files:
                with st.expander("相似片段"):
                    for relvt_cont in relvt:
                        st.write(relvt_cont)


def csv_page():
    import pandas as pd
    # import streamlit as st

    from agent.csv_agent_data import get_data_from_csv

    def create_chart(input_data, chart_type):
        df_data = pd.DataFrame(input_data["data"], columns=input_data["columns"])
        df_data.set_index(input_data["columns"][0], inplace=True)
        if chart_type == "table":
            st.table(df_data)
        elif chart_type == "bar":
            st.bar_chart(df_data)
        elif chart_type == "line":
            st.line_chart(df_data)
        elif chart_type == "scatter":
            st.line_chart(df_data)

    with st.sidebar:
        openai_key = st.text_input("输入OpenAI API Key", type="password")

    st.header("CSV智能数据分析")
    df = st.file_uploader("上传CSV文件", type=["csv"])
    if df:
        st.session_state["df"] = pd.read_csv(df)
        with st.expander("展示数据"):
            st.dataframe(st.session_state["df"])

    question = st.text_input("请根据上面提交的csv文件，进行问题提问(支持简单文本，表格，折线图，条形图，散点图)的显示")
    flag = st.button("提交问题")

    if flag:
        if 'df' in st.session_state:
            if st.session_state['df'].empty:
                st.warning("请先上传CSV文件")
                st.stop()
        else:
            st.warning("请先上传CSV文件")
            st.stop()
        if not question:
            st.warning("请输⼊问题")
            st.stop()
        if not openai_key:
            st.info("请输⼊你的Openai密钥")
            st.stop()
        with st.spinner("正在分析数据"):
            response = get_data_from_csv(openai_key, st.session_state['df'], question)
            if "answer" in response:
                st.write(response["answer"])
            if "table" in response:
                create_chart(response["table"], "table")
            if "bar" in response:
                create_chart(response["bar"], "bar")
            if "line" in response:
                create_chart(response["line"], "line")
            if "scatter" in response:
                create_chart(response["scatter"], "scatter")


def db_page():
    import sqlite3
    import tempfile

    # import streamlit as st
    from langchain_community.utilities import SQLDatabase

    from langchain_db.db_assitant import get_response

    with st.sidebar:
        openai_key = st.text_input("输入OpenAI API Key", type="password")

    st.header("SQL智能数据分析")
    db_file = st.file_uploader("上传数据库文件", type=["db"])
    if db_file:
        tfile = tempfile.NamedTemporaryFile(delete=False)
        tfile.write(db_file.getvalue())
        st.session_state["db_file"] = tfile.name
        with st.expander("展示数据"):
            conn = sqlite3.connect(st.session_state["db_file"])
            cursor = conn.cursor()
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
            tables = cursor.fetchall()
            st.write("数据库中的表：", tables)

    question = st.text_input("请根据上面提交的数据库文件，进行问题提问")
    flag = st.button("提交问题")

    if flag:
        if 'db_file' not in st.session_state:
            st.warning("请先上传数据库文件")
            st.stop()
        if not question:
            st.warning("请输入问题")
            st.stop()
        if not openai_key:
            st.info("请输入你的Openai密钥")
            st.stop()
        with st.spinner("正在分析数据"):
            db_session = st.session_state["db_file"]
            db = SQLDatabase.from_uri(f"sqlite:///{db_session}")
            response = get_response(openai_key, question, db)
            st.write(response["output"])


def ultra_ai():
    import re
    import pandas as pd
    import streamlit as st

    from ultra_ai.api.agent_config import create_agent
    from ultra_ai.api.model_config import create_llm_model

    def get_prompt_content_var(prompt_content):
        """正则表达式模式，用于匹配大括号中的内容"""
        pattern = r"\{(.*?)\}"
        matches = re.findall(pattern, prompt_content)
        table = [{"command": key, "rating": ""} for key in matches]
        return table

    def create_table(prompt_content):
        if prompt_content:
            table = get_prompt_content_var(prompt_content)
            if table:
                df = pd.DataFrame(table)
                edited_df = st.data_editor(
                    df,
                    column_config={
                        "command": "Key",
                        "rating": st.column_config.TextColumn(
                            "Value",
                            help="请输入目前key的对应value"
                        )
                    },
                    disabled=["command"],
                    hide_index=True,
                )
                return edited_df
        return None

    # st.set_page_config(layout="wide")

    # 初始化侧边栏
    with st.sidebar:
        st.title("自定义GPTs对话系统")
        apply_args = st.button("应用改变(这会重启会话)")
        st.divider()

        # 对话前提示词
        prompt = st.text_area("对话前提示词：", placeholder="可自定义提示词，变量用{var}定义")
        st.write("提示词用于对AI的回复做出一系列指令和约束。可插入表单变量，如{input}。这段提示词不会被最终用户所看到")

        # 提取提示词中的变量并创建表格
        edited_df = create_table(prompt)
        st.divider()

        # 知识库挂载
        uploaded_files = st.file_uploader("挂载企业知识库", accept_multiple_files=True, type=["pdf", "txt", "db"])
        if uploaded_files:
            st.write('你选择了以下文件：')
            for uploaded_file in uploaded_files:
                st.write(uploaded_file.name)

        st.divider()

        # 工具选择
        selected_tools = st.multiselect("请选择使用工具：", ["Python解释器", "天气工具", "数据库查询工具"])

        st.divider()

        # 开场白选择
        use_opening = st.checkbox("是否需要开场白", value=True)

    # 标题和模型配置按钮
    st.title("自定义GPTs对话系统")
    with st.popover("模型配置", help="配置OpenAI GPTs模型参数"):
        apply_model = st.button("应用改变(会重启会话)")
        api_key = st.text_input("输入密钥：", type="password")
        base_url = st.text_input("Base URL：", "https://api.aigc369.com/v1")
        model_name = st.selectbox("模型选择：", ["gpt-3.5-turbo", "gpt-4o", "gpt-4"])
        temperature = st.slider("温度选择：", 0.0, 1.0, 0.5)
        max_tokens = st.number_input("最大token选择：", min_value=1, max_value=4096, value=500)

    st.divider()

    if not api_key:
        st.info("请输⼊你的密钥")
        st.stop()

    # 初始化模型
    if apply_model or apply_args or "model" not in st.session_state:
        st.session_state["model"] = create_llm_model(api_key, base_url, model_name, temperature, max_tokens)
        print("model创建")

    if apply_model or apply_args or "final_prompt" not in st.session_state:
        variable_values = {}
        if edited_df is not None:
            for index, row in edited_df.iterrows():
                variable_values[row["command"]] = row["rating"]

        st.session_state["final_prompt"] = prompt.format(**variable_values) + "\n"

    if apply_model or apply_args or "agent" not in st.session_state:
        st.session_state["agent"] = create_agent(st.session_state["model"], selected_tools, uploaded_files,
                                                 st.session_state["final_prompt"])

    # 开场白
    if apply_model or apply_args or "messages" not in st.session_state:
        st.session_state["messages"] = [{"role": "ai", "content": "你好，我是逆蝶，你有什么想说的"}]
    if not use_opening:
        st.session_state["messages"] = []

    for message in st.session_state["messages"]:
        with st.chat_message(message["role"]):
            st.write(message["content"])

    # 对话框
    user_input = st.chat_input("我在看着你👁️")

    if user_input:
        st.session_state["messages"].append({"role": "human", "content": user_input})
        with st.chat_message("human"):
            st.write(user_input)

    if user_input:
        with st.spinner("AI正在思考🥵🥵🥵🥵🥵"):
            _input = "你的设定是：###设定开始" + st.session_state["final_prompt"] + "设定结束###" + user_input
            response = st.session_state["agent"].invoke({"input": _input})
            st.session_state["messages"].append({"role": "ai", "content": response["output"]})
            with st.chat_message("ai"):
                st.write(response["output"])


###########################

st.set_page_config(layout="wide")
with st.sidebar:
    st.title("Pizza的AI应用")

    page = st.sidebar.selectbox("点击选择其他应用",
                                ["小红书文案生成器", "视频文案生成器", "AI助手", "CSV智能数据分析", "SQL智能数据分析",
                                 "自定义GPTs"])

    st.divider()

if page == "小红书文案生成器":
    st.info("欢迎使用小红书文案生成器")
    redbook_page()
elif page == "视频文案生成器":
    st.info("欢迎使用视频文案生成器")
    video_page()
elif page == "AI助手":
    st.info("欢迎使用AI助手")
    ai_page()
elif page == "CSV智能数据分析":
    st.info("欢迎使用CSV智能数据分析")
    csv_page()
elif page == "SQL智能数据分析":
    st.info("欢迎使用SQL智能数据分析")
    db_page()
elif page == "自定义GPTs":
    st.info("欢迎使用自定义GPTs")
    ultra_ai()
